import gzip
import json
import logging
import os
import shutil
import tarfile
import threading
import zipfile

from myapp.app import db
from myapp.models.model_dataset import Dataset
from myapp.utils.storage import storage_mgr
from myapp.utils.unpack_check import check_dataset_path


def unpack_succeed(dataset_id, unpack_path):
    m2 = db.session.query(Dataset).filter_by(id=int(dataset_id)).first()
    m2.status = 'succeed'
    m2.unpack_path = unpack_path
    m2.dataset_path = unpack_path
    db.session.commit()


def unpack_fail(dataset_id, fail_msg):
    m2 = db.session.query(Dataset).filter_by(id=int(dataset_id)).first()
    m2.status = 'failed'
    m2.expand = json.dumps({'fail_msg': fail_msg})
    db.session.commit()


# 异步解压数据集
def unpack_dataset(dataset_id, callback=None, callback_args=()):
    threading.Thread(
        target=unpack_dataset_sync, args=(dataset_id, callback, callback_args)
    ).start()


# 同步解压数据集
def unpack_dataset_sync(dataset_id, callback=None, callback_args=()):
    model_d = db.session.query(Dataset).filter_by(id=int(dataset_id)).first()
    # 异常
    if model_d is None:
        if callback:
            callback(callback_args, is_success=False, fail_msg='数据集不存在' + str(dataset_id))
        return

    # 解压
    unpack_path = model_d.gen_unpack_path()
    fail_msg = do_unpack(
        model_d.origin_path, unpack_path, model_d.label_type, model_d.labeled == 1
    )
    is_success = True if fail_msg is None else False

    # 更新dataset表
    if is_success:
        unpack_succeed(dataset_id, unpack_path)
    else:
        unpack_fail(dataset_id, fail_msg)

    # 回调
    if callback:
        callback(callback_args, is_success=is_success, fail_msg=fail_msg)


# 异步解压目录
# def unpack_package_async(origin_path, unpack_path):
#     threading.Thread(target=do_unpack, args=(origin_path, unpack_path)).start()


def do_unpack(origin_path, unpack_path, label_type, need_check_dir=False):
    # bucket = cube_bucket

    if unpack_path[-1] == '/':
        unpack_path = unpack_path[:-1]

    # 拉取到本地
    local_origin = './' + os.path.basename(origin_path)
    logging.info('开始下载到：local_origin: ' + local_origin)
    try:
        logging.info(f'开始下载{origin_path}到{local_origin}')
        storage_mgr.download_file(origin_path, local_origin)
    except Exception:
        import traceback

        logging.info(traceback.format_exc())
        return '文件下载失败，请重新创建任务'

    # 解压
    ext = os.path.splitext(local_origin)[-1].lower()
    logging.info('下载完成,开始unpack, 格式为:' + ext)

    try:
        if ext == '.tar':
            unpack = un_tar(local_origin)
        elif ext == '.gz':
            unpack = un_gz(local_origin)
        elif ext == '.zip':
            res = un_zip2(local_origin)
            logging.info('UnzipResult: ' + res[1])
            if res[0] < 0:
                logging.info('解压失败', res[1], origin_path, unpack_path)
                return res[1]
            unpack = res[1]
        else:
            # set_upload_fail(model_dataset.id, "解压失败：不支持的压缩格式")
            logging.info('解压失败', origin_path, unpack_path)
            return '文件解压失败，请重新创建任务'
    except Exception as e:
        logging.exception(e)
        return '文件解压失败，请重新创建任务'

    # 校验
    if need_check_dir:
        check_result = check_dataset_path(unpack, label_type)
        if check_result is not None:
            # 清理
            logging.info(f'格式校验失败{check_result}，开始清理')
            shutil.rmtree(local_origin.replace('.zip', ''))
            os.remove(local_origin)
            logging.info('清理完成')
            return check_result

    # 回传
    logging.info('localUnpackPath: ' + unpack)
    logging.info('unpack成功，开始上传到obs_load_to: ' + unpack_path)
    try:
        storage_mgr.upload_dir(unpack, unpack_path)
    except Exception:
        logging.info('上传失败', origin_path, unpack_path)
        return '文件上传失败，请重新创建任务'

    # 清理
    logging.info('上传成功，开始清理')
    # shutil.rmtree(unpack)
    shutil.rmtree(local_origin.replace('.zip', ''))
    os.remove(local_origin)
    logging.info('清理完成')
    return None


def is_compressed(p):
    if p is None or len(p) < 1:
        # 为空判定为False
        return False
    ext = os.path.splitext(p)[-1].lower()
    # 列出压缩包格式
    if ext in ['.zip', '.tar', '.gz']:
        return True
    else:
        return False


def is_alluxio():
    st_type = os.getenv('STORAGE_MEDIA')
    return st_type == 'MINIO' or st_type == 'minio' or st_type == 'arm' or st_type == 'ARM'


def un_tar(file_name):
    tar = tarfile.open(file_name)
    t = file_name.replace('.tar', '') + '_files/'
    names = tar.getnames()
    if os.path.isdir(t):
        pass
    else:
        os.mkdir(t)
    for name in names:
        tar.extract(name, t)
    tar.close()
    return t


def un_gz(file_name):
    f_name = file_name.replace('.gz', '') + '_files/'
    g_file = gzip.GzipFile(file_name)
    open(f_name, 'w+').write(g_file.read())
    g_file.close()
    return f_name


# deprecated
def un_zip(file_name):
    zip_file = zipfile.ZipFile(file_name)
    dir_filename = file_name.replace('.zip', '')
    if os.path.isdir(dir_filename):
        pass
    else:
        os.mkdir(dir_filename)
    from_path = dir_filename + '_o'
    for names in zip_file.namelist():
        if names.find('__MACOSX') != -1 or names.find('DS_Store') != -1:
            continue
        zip_file.extract(names, from_path)

        # dcode = 'utf-8'
        # names = str_decode(str_encode(names))
        to_name = dir_filename + '/' + str_decode(str_encode(names))
        # try:
        #     to_name = dir_filename + '/' + names.encode('cp437').decode(dcode)
        # except Exception as e:
        #     logging.info("NOTUTF8:" + str(e))
        #     dcode = 'gbk'
        #     to_name = dir_filename + '/' + names.encode('cp437').decode(dcode)
        # os.rename(from_path + '/' + names, to_name)

        if os.path.isfile(from_path + '/' + names):
            shutil.copy(from_path + '/' + names, to_name)
        else:
            if not os.path.isdir(to_name):
                os.makedirs(to_name)
    shutil.rmtree(from_path)
    zip_file.close()

    dirs_under_filename = os.listdir(dir_filename)
    logging.info('dirs_under_filename COUNT: ' + str(len(dirs_under_filename)))
    if len(dirs_under_filename) == 1:
        return dir_filename + '/' + dirs_under_filename[0]
    else:
        return dir_filename


# UTF_8_BOM = b'\xef\xbb\xbf'


def un_zip2(filepath):
    try:
        # encoding = 'utf-8'
        origin_unzip_pwd = filepath.replace('.zip', '') + '_origin'
        unzip_ok_pwd = filepath.replace('.zip', '')
        if not os.path.isdir(origin_unzip_pwd):
            os.mkdir(origin_unzip_pwd)
        if not os.path.exists(unzip_ok_pwd):
            os.mkdir(unzip_ok_pwd)
        logging.info('origin_unzip_pwd: %s' % origin_unzip_pwd)
        # dst_dir = try_encode(dst_dir, encoding)
        zip_file = zipfile.ZipFile(filepath)
        for names in zip_file.namelist():
            if names.find('__MACOSX') != -1 or names.find('DS_Store') != -1:
                continue
            # logging.info(try_encode(str_decode(names.encode('cp437')), encoding))
            unzip_file = zip_file.extract(names, origin_unzip_pwd)
            try:
                rename_file = os.path.join(unzip_ok_pwd, str_decode(str_encode(names)))
            except Exception as e:
                logging.info('NOTUTF8:' + str(e))
                rename_file = os.path.join(unzip_ok_pwd, names)
            if rename_file.endswith('/'):
                os.makedirs(rename_file, exist_ok=True)
            else:
                # 遍历的结果可能不包含目录，先创建目录
                dirname = os.path.dirname(rename_file)
                os.makedirs(dirname, exist_ok=True)
                os.rename(unzip_file, rename_file)
            logging.info('origin unzip file %s', unzip_file)
            logging.info('origin rename_file file %s', rename_file)
        zip_file.close()
        shutil.rmtree(origin_unzip_pwd)

        dirs_under_filename = os.listdir(unzip_ok_pwd)
        inner_path = os.path.join(unzip_ok_pwd, dirs_under_filename[0])
        # logging.info("len(dirs_under_filename):" + str(len(dirs_under_filename)))
        # logging.info("inner_path:" + inner_path)
        for d in dirs_under_filename:
            logging.info('dirs_under_filename:' + d)
        # return 0, inner_path
        if len(dirs_under_filename) == 1 and os.path.isdir(inner_path):
            logging.info('ChoseA: ' + inner_path)
            return 0, inner_path
        else:
            logging.info('ChoseB: ' + unzip_ok_pwd)
            return 0, unzip_ok_pwd
    except Exception as err:
        logging.exception(err)
        logging.info('[un_zip]: Exception.ERROR: ', err)
        return -1, '您上传的文件路径或文件名包含不支持的特殊字符，请修改后重新上传'


def str_encode(string):
    while True:
        dec = try_encode(string, 'cp437')
        if dec is not None:
            break
        dec = try_encode(string, 'utf-8')
        if dec is not None:
            break
        dec = try_encode(string, 'ascii')
        if dec is not None:
            break
        dec = try_encode(string, 'GB2312')
        if dec is not None:
            break
        dec = try_encode(string, 'GBK')
        if dec is not None:
            break
        dec = try_encode(string, 'Big5')
        if dec is not None:
            break
        logging.info('[try_encode]: unknown encoding')
        dec = None
        break
    return dec


def try_encode(s, encoding='utf-8'):
    if s is None:
        logging.info('[tryEncode]: input param None!')
        return s
    try:
        return s.encode(encoding)
    except UnicodeEncodeError:
        # logging.info(err)
        return None


def try_decode(s, decoding='utf-8'):
    try:
        return s.decode(decoding)
    except UnicodeDecodeError:
        # logging.info(err)
        return None


def str_decode(string):
    while True:
        dec = try_decode(string, 'utf-8')
        if dec is not None:
            break
        dec = try_decode(string, 'GB2312')
        if dec is not None:
            break
        dec = try_decode(string, 'GBK')
        if dec is not None:
            break
        dec = try_decode(string, 'ascii')
        if dec is not None:
            break
        dec = try_decode(string, 'Big5')
        if dec is not None:
            break
        logging.info('[str_decode]: unknown encoding')
        dec = None
        break

    return dec
